""" Merge the contract and well datasets with the mapping between them.
Also, make the rig specification from the max. water depth in the
mapping.

"""
import pandas as pd
import utils

#%% SET UP PARAMETERS -------------------------------------------------------------------
bins_spec = pd.IntervalIndex.from_tuples([(0, 200.5), (200.6, 299.4), (299.5, 500)])
workover_rigs = [
    "dolphin 110",
    "hercules 211",
    "ocean summit",
    "p-54",
    "paragon m824",
    "spartan 151",
    "dolphin 109",
    "well services rig 53",
    "d-105",
    "dolphin 106",
    "p-50",
    "pool ranger v",
    "pool ranger v",
    "pool ranger vii",
    "dolphin 106",  # nabors
    "p-50",  # nabors
    "pool ranger v",  # nabors
]
# Pride wisconsin/pacific 12 does not seem to be in the well data

#%% IMPORT MAPPING; CONSTRUCT RIG SPECIFICATION -----------------------------------------
map = (
    pd.read_excel('./data_py/external/rig_map_VALIDATE.xls', sheet_name='Sheet1')
    .dropna(subset=['id_in_well'])
    .assign(spec=lambda x:
        pd.cut(x['max_wd'], bins=bins_spec, labels=['low', 'mid', 'high'])
    )
)

#%% MERGE WELLS WITH MAPPING ------------------------------------------------------------
wells_with_map = (
    pd.read_csv('./data_py/temp/03_make_mri/wells_with_mri.csv', index_col=[0])
    .dropna(subset=['rig_id_num'])
    .assign(rig_id_num=lambda x: x['rig_id_num'].astype(int))
    .replace({4818: 4807})  # Replace multiple names for noble tom jobe
    .merge(
        map,
        left_on='rig_id_num',
        right_on='id_in_well',
        validate='m:1',
        how='left',
        indicator=True
    )
    .assign(rig_name=lambda x: x['rig_name'].str.lower())
    .query('name_in_ihs != @workover_rigs')
    .pipe(
        utils.to_csv_return,
        path='./data_py/temp/05_merge_with_map/wells_with_map_incl_unmerged.csv'
    )
    .query('_merge == "both"')
    .drop('_merge', axis=1)
    .pipe(
        utils.to_csv_return,
        path='./data_py/temp/05_merge_with_map/wells_with_map.csv'
    )
)

#%% MERGE CONTRACTS WITH MAPPING --------------------------------------------------------
contracts_with_map = (
    pd.read_csv('./data_py/temp/04_clean_contracts/contracts_processed.csv', index_col=[0])
    .merge(
        map,
        left_on='rig_name',
        right_on='name_in_ihs',
        validate='m:1',
        how='left',
        indicator=True
    )
    .query('name_in_ihs != @workover_rigs')
    .pipe(
        utils.to_csv_return,
        path='./data_py/temp/05_merge_with_map/contracts_with_map_incl_unmerged.csv'
    )
    .query('_merge == "both"')
    .drop('_merge', axis=1)
    .pipe(
        utils.to_csv_return,
        path='./data_py/temp/05_merge_with_map/contracts_with_map.csv'
    )
)
